import tensorflow as tf
from tensorflow.keras import layers
import numpy as np

# class MyModel(Model):
#     def __init__(self, *args, **kwargs):
#         super().__init__(*args, **kwargs)
#         self.d1 = layers.Dense(8, activation='relu')
#         self.d2 = layers.Dense(1, activation='sigmoid')
#
#     @tf.function
#     def call(self, inputs, training=None, mask=None):
#         x = self.d1(inputs)
#         y = self.d2(x)
#         return y

# 取得模型函数
'''
传入的my_mat_1和my_mat_2是通过GA优化的权值矩阵, 将其作为网络训练的初始值
'''


def get_model(my_mat_1, my_mat_2):
    model = tf.keras.Sequential([
        layers.Dense(56, activation='relu', kernel_initializer=my_mat_1),
        layers.Dense(1, activation='sigmoid', kernel_initializer=my_mat_2)
    ])

    return model
